import os
import requests
import time
from tavily import TavilyClient
from dotenv import load_dotenv

# Load environment variables from .env file
load_dotenv()

# Get Tavily API key from environment or use the provided one
TAVILY_API_KEY = os.getenv("TAVILY_API_KEY", "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")

# Create images directory if it doesn't exist
os.makedirs("images", exist_ok=True)

# Initialize Tavily client
client = TavilyClient(TAVILY_API_KEY)

def download_images(query, count, prefix):
    """
    Download images using Tavily API
    
    Args:
        query: Search query
        count: Number of images to download
        prefix: Prefix for saved image filenames
    """
    print(f"Downloading {count} images for '{query}'...")
    
    try:
        # Get image results using Tavily search with include_images parameter
        response = client.search(query=query, include_images=True, max_results=count)
        
        # Extract images from the response
        images = response.get("images", [])
        
        if not images:
            print(f"No images found for '{query}'")
            return
        
        for i, image_url in enumerate(images):
            if i >= count:
                break
                
            # Download the image
            try:
                img_response = requests.get(image_url, timeout=10)
                img_response.raise_for_status()
                
                # Save the image
                filename = f"images/{prefix}_{i+1}.jpg"
                with open(filename, "wb") as f:
                    f.write(img_response.content)
                
                print(f"Downloaded: {filename}")
                
                # Add a small delay to avoid rate limiting
                time.sleep(0.5)
                
            except Exception as e:
                print(f"Error downloading image {i+1}: {e}")
    
    except Exception as e:
        print(f"Error searching for images: {e}")

# Main execution
if __name__ == "__main__":
    # Download images of people wearing masks
    download_images("people wearing face masks clear photo", 10, "mask")
    
    # Download images of people without masks
    download_images("people face without mask clear photo", 10, "no_mask")
    
    print("Download completed!")
